import argparse
import os
import logging
import time
from pathlib import Path
import json
from tqdm import tqdm

from docling_core.types.doc import ImageRefMode, PictureItem, TableItem, TextItem
from docling.datamodel.base_models import FigureElement, InputFormat, Table
from docling.datamodel.pipeline_options import PdfPipelineOptions, AcceleratorOptions, AcceleratorDevice, RapidOcrOptions
from docling.document_converter import DocumentConverter, PdfFormatOption, ConversionResult
from openai import OpenAI

# 设置日志
_log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

class DocumentProcessor:
    def __init__(self):
        self.IMAGE_RESOLUTION_SCALE = 5.0
        self.artifacts_path = r"./models/docling-models/"
        self.client = OpenAI(
            api_key="sk-ae6d05c39cd94735840397a0f7f51fb8",
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
        
    def process_document(self, input_pdf_path: str, output_dir_name: str = "processed_docs") -> ConversionResult:
        """处理PDF文档，转换为多种格式并提取内容"""
        output_dir = Path(output_dir_name)
        input_doc_path = input_pdf_path
        
        # 设置OCR模型路径
        download_path = r"./models/RapidOCR/onnx/"
        
        # 构建OCR选项
        det_model_path = os.path.join(download_path, "PP-OCRv4/", "ch_PP-OCRv4_det_infer.onnx")
        rec_model_path = os.path.join(download_path, "PP-OCRv4/", "ch_PP-OCRv4_rec_server_infer.onnx")
        cls_model_path = os.path.join(download_path, "PP-OCRv4/", "ch_ppocr_mobile_v2.0_cls_infer.onnx")

        _log.info(f"det_model_path: {det_model_path}")
        _log.info(f"rec_model_path: {rec_model_path}")
        _log.info(f"cls_model_path: {cls_model_path}")

        ocr_options = RapidOcrOptions(
            det_model_path=det_model_path,
            rec_model_path=rec_model_path,
            cls_model_path=cls_model_path,
        )

        pipeline_options = PdfPipelineOptions(
            ocr_options=ocr_options,
            artifacts_path=self.artifacts_path
        )
        
        # 配置处理选项
        pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.CUDA)
        pipeline_options.images_scale = self.IMAGE_RESOLUTION_SCALE
        pipeline_options.generate_page_images = True
        pipeline_options.generate_picture_images = True

        # 创建文档转换器
        doc_converter = DocumentConverter(
            format_options={
                InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)
            }
        )

        _log.info("开始文档转换...")
        start_time = time.time()
        
        # 执行转换
        conv_res = doc_converter.convert(input_doc_path)
        doc_filename = conv_res.input.file.stem
        output_dir = Path(os.path.join(output_dir, doc_filename))
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 保存页面图片
        for page_no, page in tqdm(conv_res.document.pages.items()):
            page_no = page.page_no
            page_image_filename = output_dir/f"{doc_filename}-{page_no}.png"
            with page_image_filename.open("wb") as fp:
                page.image.pil_image.save(fp, format="PNG")

        # 保存表格和图片
        table_counter = 0
        picture_counter = 0
        for element, _level in tqdm(conv_res.document.iterate_items()):
            if isinstance(element, TableItem):
                table_counter += 1
                element_image_filename = (output_dir / f"{doc_filename}-table-{table_counter}.png")
                with element_image_filename.open("wb") as fp:
                    element.get_image(conv_res.document).save(fp, "PNG")

            if isinstance(element, PictureItem):
                picture_counter += 1
                element_image_filename = (output_dir / f"{doc_filename}-picture-{picture_counter}.png")
                with element_image_filename.open("wb") as fp:
                    element.get_image(conv_res.document).save(fp, "PNG")

        base_filename = output_dir / doc_filename
        conv_res.document.save_as_markdown(f"{base_filename}-with-image-refs.md", image_mode=ImageRefMode.REFERENCED)
        conv_res.document.save_as_html(f"{base_filename}-with-image-refs.html", image_mode=ImageRefMode.REFERENCED)
        conv_res.document.save_as_json(f"{base_filename}-with-image-refs.json", image_mode=ImageRefMode.REFERENCED)
        
        end_time = time.time()
        execution_time = end_time - start_time
        _log.info(f"文档处理完成，耗时: {execution_time:.2f}秒")
        
        return conv_res

    def extract_use_cases(self, md_path: str):
        """从Markdown文档中提取用例信息"""
        _log.info("开始提取用例信息...")
        
        msg = []
        file_object = self.client.files.create(file=Path(md_path), purpose="file-extract")
        msg.append({'role': 'system', 'content': f'fileid://{file_object.id}'})
        msg.append({
            'role': 'user', 
            'content': '''
            请从文档中提取所有软件需求信息,以纯文本json列表的方式返回，要求如下：

            不要包含任何解释性文字和格式化Markdown语法(无需用如```json  ```包围)

            1. 需求分类：
            - 功能需求：系统必须实现的具体功能
            - 性能需求：系统在性能指标上的要求
            - 安全需求：系统在安全性方面的要求
            - 其他需求：不属于上述分类的需求

            2. 每个需求包含以下字段：
            - "Name": 需求名称/标题（简洁明确）
            - "ID": 需求唯一标识（如文档中有）
            - "Description": 需求详细描述
            - "Type": 需求类型（功能需求/性能需求/非功能需求）
            - "Priority": 优先级（高/中/低，如文档中有）
            - "Source": 需求来源（如用户故事、法规要求等）

            3. 提取要求：
            - 保持原始需求的表述，不要改写或概括
            - 忽略文档中的目录、标题等非需求内容
            - 如果需求有测试相关描述，请保留在Description中
            - 确保需求之间的独立性

            4. 输出格式：
            - 严格的JSON数组
            - 每个对象对应一个需求
            - 不要包含任何解释性文字和Markdown语法

            示例：
            [
                {
                    "Name": "用户登录功能",
                    "ID": "REQ-001",
                    "Description": "系统应提供用户名密码登录方式，支持5次错误尝试后锁定账户30分钟",
                    "Type": "功能需求",
                    "Priority": "高",
                    "Source": "用户故事US-101"
                },
                {
                    "Name": "系统响应时间",
                    "ID": "REQ-042",
                    "Description": "在100并发用户下，主要页面加载时间应小于2秒",
                    "Type": "性能需求",
                    "Priority": "中",
                    "Source": "性能指标要求"
                }
            ]

            请开始提取文档中的需求信息：
            '''
        })

        completion = self.client.chat.completions.create(
            model="qwen-long",
            messages=msg,
            temperature=0.3
        )
        
        result = json.loads(completion.model_dump_json())
        return result

    def extract_test_results(self, md_path: str):
        """从测试报告中提取测试结果"""
        _log.info("开始提取测试结果...")
        
        msg = []
        file_object = self.client.files.create(file=Path(md_path), purpose="file-extract")
        msg.append({'role': 'system', 'content': f'fileid://{file_object.id}'})
        msg.append({
            'role': 'user', 
            'content': '''
            请从测试报告中提取测试用例及其结果,以纯文本的方式返回，要求如下：

            不要包含任何解释性文字和格式化Markdown语法(如```json ```)

            1. 测试用例结构：
            - "Name": 测试用例名称/标题
            - "ID": 测试用例编号
            - "RequirementID": 对应的需求 ID （如可识别）
            - "TestObjective": 测试的需求类型（功能需求/性能需求/非功能需求）
            - "TestMethod": 测试方法描述（步骤、工具、环境等）
            - "TestResult": 测试结果描述
            - "Status": 测试状态（通过/失败/阻塞/未执行）
            - "Evidence": 测试证据（截图、日志等引用）
            - "DefectID": 关联的缺陷ID（如有）

            2. 提取要求：
            - 保持原始测试描述的完整性
            - 尽可能关联到具体需求
            - 区分测试方法和测试结果
            - 保留测试数据和环境信息
            - 忽略文档中的目录、统计信息等非用例内容

            3. 输出格式：
            - 严格的JSON数组
            - 每个对象对应一个测试用例
            - 不要包含任何解释性文字和Markdown语法

            示例：
            [
                {
                    "Name": "用户登录功能测试",
                    "ID": "TC-101",
                    "RequirementID": "REQ-001",
                    "TestObjective": "验证用户登录功能是否符合要求",
                    "TestMethod": "1. 输入正确用户名密码\n2. 输入错误密码(5次)\n3. 验证账户锁定\n工具：Selenium自动化测试",
                    "TestResult": "1. 成功登录\n2. 第5次错误后账户锁定\n3. 30分钟后自动解锁",
                    "Status": "通过",
                    "Evidence": "截图见附件1",
                    "DefectID": ""
                },
                {
                    "Name": "系统负载测试",
                    "ID": "TC-205",
                    "RequirementID": "REQ-042",
                    "TestObjective": "验证系统在100并发下的响应时间",
                    "TestMethod": "使用JMeter模拟100并发用户，持续运行30分钟",
                    "TestResult": "平均响应时间1.8秒，满足<2秒要求",
                    "Status": "通过",
                    "Evidence": "JMeter报告见附件3",
                    "DefectID": ""
                }
            ]

            请开始提取测试报告中的测试用例信息：
            '''
        })

        completion = self.client.chat.completions.create(
            model="qwen-long",
            messages=msg,
            temperature=0.2  # 更低温度保证结果稳定性
        )
        
        result = json.loads(completion.model_dump_json())
        return result

    def save_as_readable_text(self, data: dict, output_path: Path, title: str = "文档解析结果"):
        """将解析结果保存为易读的文本格式"""
        with open(output_path, 'w', encoding='utf-8') as f:
            # 写入标题和生成时间
            f.write(f"=== {title} ===\n")
            f.write(f"生成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            
            # 写入需求信息
            if 'use_cases' in data and data['use_cases']:
                f.write("=== 需求信息 ===\n\n")
                use_cases = self._parse_api_response(data['use_cases'])
                for i, case in enumerate(use_cases, 1):
                    f.write(f"需求 #{i}\n")
                    f.write(f"名称: {case.get('Name', '无')}\n")
                    f.write(f"ID: {case.get('ID', '无')}\n")
                    f.write(f"类型: {case.get('Type', '无')}\n")
                    f.write(f"优先级: {case.get('Priority', '无')}\n")
                    f.write(f"来源: {case.get('Source', '无')}\n")
                    f.write(f"描述:\n{case.get('Description', '无')}\n")
                    f.write("\n" + "-"*50 + "\n\n")
            
            # 写入测试结果
            if 'test_results' in data and data['test_results']:
                f.write("\n=== 测试结果 ===\n\n")
                test_results = self._parse_api_response(data['test_results'])
                for i, test in enumerate(test_results, 1):
                    f.write(f"测试用例 #{i}\n")
                    f.write(f"名称: {test.get('Name', '无')}\n")
                    f.write(f"ID: {test.get('ID', '无')}\n")
                    f.write(f"关联需求ID: {test.get('RequirementID', '无')}\n")
                    f.write(f"测试目的: {test.get('TestObjective', '无')}\n")
                    f.write(f"测试方法:\n{test.get('TestMethod', '无')}\n")
                    f.write(f"测试结果:\n{test.get('TestResult', '无')}\n")
                    f.write(f"状态: {test.get('Status', '无')}\n")
                    f.write(f"证据: {test.get('Evidence', '无')}\n")
                    f.write(f"缺陷ID: {test.get('DefectID', '无')}\n")
                    f.write("\n" + "-"*50 + "\n\n")
            
            f.write("\n=== 解析完成 ===\n")

    def _parse_api_response(self, data):
        """解析API返回的数据结构"""
        if isinstance(data, dict) and 'choices' in data:
            try:
                return json.loads(data['choices'][0]['message']['content'])
            except:
                return []
        elif isinstance(data, (list, dict)):
            return data
        else:
            return []

    def full_processing_pipeline(self, requirement_pdf: str, test_report_pdf: str):
        """完整的文档处理流程"""
        doc_name = Path(requirement_pdf).stem
        output_dir = Path(f"processed_docs/{doc_name}")
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 1. 处理需求文档
        _log.info("\n===== 开始处理需求文档 =====")
        req_conv_result = self.process_document(requirement_pdf)
        req_md_path = output_dir / f"{doc_name}-with-image-refs.md"
        use_cases = self.extract_use_cases(str(req_md_path))

        doc_name = Path(test_report_pdf).stem
        output_dir = Path(f"processed_docs/{doc_name}")
        
        # 2. 处理测试报告
        _log.info("\n===== 开始处理测试报告 =====")
        test_conv_result = self.process_document(test_report_pdf)
        test_md_path = output_dir / f"{Path(test_report_pdf).stem}-with-image-refs.md"
        test_results = self.extract_test_results(str(test_md_path))
        
        # 3. 保存结果
        results = {
            "use_cases": use_cases,
            "test_results": test_results
        }
        
        # 保存为易读的文本格式
        readable_path = output_dir / "extracted_results.txt"
        self.save_as_readable_text(results, readable_path, f"{doc_name} 解析结果")
        
        _log.info(f"结果已保存为易读格式: {readable_path}")
        return results

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='文档处理工具')
    parser.add_argument('-r', '--requirement', type=str, required=True, 
                        help='需求文档PDF文件路径')
    parser.add_argument('-t', '--test', type=str, required=True,
                        help='测试报告PDF文件路径')
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_arguments()
    processor = DocumentProcessor()
    
    # 完整处理流程
    results = processor.full_processing_pipeline(args.requirement, args.test)
    print("\n处理完成，结果已保存到processed_docs目录")